/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.junit.BeforeClass;
import org.junit.Test;
import static org.junit.Assert.*;

import javax.net.ssl.*;
import java.io.*;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.security.*;
import java.security.cert.X509Certificate;

import org.openhitls.crypto.jce.provider.HiTls4jProvider;

/**
 * Test real HTTPS connection using HiTLS provider.
 */
public class HttpsConnectionTest {

    private static final String PROVIDER_NAME = HiTls4jProvider.PROVIDER_NAME;

    @BeforeClass
    public static void loadNativeLibraries() {
        String libraryPath = new File(System.getProperty("user.dir") + "/target/native/libhitls_crypto_jni.so").getAbsolutePath();
        System.load(libraryPath);
        Security.addProvider(new HiTls4jProvider());
    }

    @Test
    public void testHttpsConnectionToPqcDemo() throws Exception {
        System.out.println("Testing HTTPS connection to pqc.ssldemo.cn with CA certificate verification...");

        // Create SSLContext with system CA certificate verification
        // By passing null for TrustManager, we use system default CA certificates
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, new SecureRandom());
        System.out.println("Using system CA certificates for server certificate verification");

        // Create SSLEngine for client mode
        SSLEngine engine = ctx.createSSLEngine("pqc.ssldemo.cn", 443);
        engine.setUseClientMode(true);

        // Get buffer sizes
        SSLSession session = engine.getSession();
        int appBufferSize = session.getApplicationBufferSize();
        int netBufferSize = session.getPacketBufferSize();

        System.out.println("App buffer size: " + appBufferSize);
        System.out.println("Net buffer size: " + netBufferSize);

        // Allocate buffers
        ByteBuffer appOut = ByteBuffer.allocate(appBufferSize);
        ByteBuffer appIn = ByteBuffer.allocate(appBufferSize);
        ByteBuffer netOut = ByteBuffer.allocate(netBufferSize);
        ByteBuffer netIn = ByteBuffer.allocate(netBufferSize);

        // Connect to PQC Demo server
        try (Socket socket = new Socket("pqc.ssldemo.cn", 443)) {
            socket.setSoTimeout(30000); // 30 second timeout
            InputStream in = socket.getInputStream();
            OutputStream out = socket.getOutputStream();

            // Begin handshake
            engine.beginHandshake();
            System.out.println("Starting TLS handshake...");

            // Handshake loop
            SSLEngineResult.HandshakeStatus hsStatus = engine.getHandshakeStatus();
            int handshakeLoops = 0;
            final int MAX_LOOPS = 100;

            while (hsStatus != SSLEngineResult.HandshakeStatus.FINISHED &&
                   hsStatus != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {

                if (++handshakeLoops > MAX_LOOPS) {
                    fail("Handshake took too many iterations");
                }

                System.out.println("Handshake status: " + hsStatus);

                switch (hsStatus) {
                    case NEED_WRAP:
                        netOut.clear();
                        SSLEngineResult wrapResult = engine.wrap(appOut, netOut);
                        System.out.println("Wrap result: " + wrapResult);
                        hsStatus = wrapResult.getHandshakeStatus();

                        if (netOut.position() > 0) {
                            netOut.flip();
                            byte[] data = new byte[netOut.remaining()];
                            netOut.get(data);
                            out.write(data);
                            out.flush();
                            System.out.println("Sent " + data.length + " bytes to server");
                        }
                        break;

                    case NEED_UNWRAP:
                        // Read from network with retry for slow connections
                        byte[] readBuf = new byte[netBufferSize];
                        int bytesRead = -1;
                        int readAttempts = 0;
                        while (bytesRead <= 0 && readAttempts < 3) {
                            try {
                                bytesRead = in.read(readBuf);
                            } catch (java.net.SocketTimeoutException e) {
                                readAttempts++;
                                System.out.println("Read timeout, attempt " + readAttempts);
                                continue;
                            }
                            if (bytesRead < 0) {
                                fail("Connection closed during handshake");
                            }
                        }
                        if (bytesRead <= 0) {
                            fail("Failed to read data after retries");
                        }
                        System.out.println("Received " + bytesRead + " bytes from server");

                        netIn.clear();
                        netIn.put(readBuf, 0, bytesRead);
                        netIn.flip();

                        SSLEngineResult unwrapResult = engine.unwrap(netIn, appIn);
                        System.out.println("Unwrap result: " + unwrapResult);
                        hsStatus = unwrapResult.getHandshakeStatus();
                        break;

                    case NEED_TASK:
                        Runnable task;
                        while ((task = engine.getDelegatedTask()) != null) {
                            task.run();
                        }
                        hsStatus = engine.getHandshakeStatus();
                        break;

                    default:
                        break;
                }
            }

            System.out.println("TLS handshake completed!");
            System.out.println("Protocol: " + engine.getSession().getProtocol());
            System.out.println("Cipher suite: " + engine.getSession().getCipherSuite());

            // Print key exchange group info
            if (engine instanceof HiTlsSSLEngine) {
                HiTlsSSLEngine hiEngine = (HiTlsSSLEngine) engine;
                System.out.println("Key exchange group: " + hiEngine.getNegotiatedGroupName());
            }

            // Send HTTP request
            String httpRequest = "GET / HTTP/1.1\r\nHost: pqc.ssldemo.cn\r\nConnection: close\r\n\r\n";
            appOut.clear();
            appOut.put(httpRequest.getBytes());
            appOut.flip();

            netOut.clear();
            SSLEngineResult wrapResult = engine.wrap(appOut, netOut);
            System.out.println("HTTP request wrap: " + wrapResult);

            if (netOut.position() > 0) {
                netOut.flip();
                byte[] data = new byte[netOut.remaining()];
                netOut.get(data);
                out.write(data);
                out.flush();
                System.out.println("Sent HTTP request (" + data.length + " bytes)");
            }

            // Read HTTP response
            StringBuilder response = new StringBuilder();
            boolean readingResponse = true;
            int responseLoops = 0;

            while (readingResponse && responseLoops++ < 50) {
                byte[] readBuf = new byte[netBufferSize];
                int bytesRead;
                try {
                    bytesRead = in.read(readBuf);
                } catch (Exception e) {
                    break;
                }

                if (bytesRead < 0) {
                    break;
                }

                netIn.clear();
                netIn.put(readBuf, 0, bytesRead);
                netIn.flip();

                appIn.clear();
                SSLEngineResult unwrapResult = engine.unwrap(netIn, appIn);
                System.out.println("Response unwrap: " + unwrapResult);

                if (appIn.position() > 0) {
                    appIn.flip();
                    byte[] appData = new byte[appIn.remaining()];
                    appIn.get(appData);
                    response.append(new String(appData));

                    // Check if we have enough of the response
                    if (response.length() > 500) {
                        readingResponse = false;
                    }
                }

                if (unwrapResult.getStatus() == SSLEngineResult.Status.CLOSED) {
                    readingResponse = false;
                }
            }

            System.out.println("\n=== HTTP Response (first 500 chars) ===");
            String responseStr = response.toString();
            System.out.println(responseStr.substring(0, Math.min(500, responseStr.length())));

            // Verify we got an HTTP response
            assertTrue("Should receive HTTP response", responseStr.startsWith("HTTP/"));
            assertTrue("Should receive 200 OK or redirect",
                responseStr.contains("200") || responseStr.contains("301") || responseStr.contains("302"));

            System.out.println("\nHTTPS connection test PASSED!");
        }
    }
}
